#include <windows.h>
#include <shellapi.h>
#include <sddl.h>
#include "ElevatePrivileges.h"

// 检查当前是否以管理员权限运行
__declspec(dllexport) BOOL IsRunAsAdmin()
{
    BOOL fIsRunAsAdmin = FALSE;
    PSID pAdministratorsGroup = NULL;
    
    // 创建管理员组的SID
    SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
    if (!AllocateAndInitializeSid(
        &NtAuthority,
        2,
        SECURITY_BUILTIN_DOMAIN_RID,
        DOMAIN_ALIAS_RID_ADMINS,
        0, 0, 0, 0, 0, 0,
        &pAdministratorsGroup))
    {
        return FALSE;
    }

    // 检查当前进程令牌是否属于管理员组
    if (!CheckTokenMembership(NULL, pAdministratorsGroup, &fIsRunAsAdmin))
    {
        fIsRunAsAdmin = FALSE;
    }

    // 释放SID
    if (pAdministratorsGroup)
    {
        FreeSid(pAdministratorsGroup);
    }

    return fIsRunAsAdmin;
}

// 以管理员权限重新启动程序
__declspec(dllexport) BOOL RunAsAdmin(
    LPCWSTR lpApplicationPath, 
    int nShowCmd, 
    LPCWSTR lpParameters, 
    LPCWSTR lpDirectory)
{
    SHELLEXECUTEINFOW sei = { sizeof(sei) };
    
    sei.lpVerb = L"runas"; // 请求提升权限
    sei.lpFile = lpApplicationPath;
    sei.hwnd = NULL;
    sei.nShow = nShowCmd;
    sei.lpParameters = lpParameters;
    sei.lpDirectory = lpDirectory;
    
    if (!ShellExecuteExW(&sei))
    {
        DWORD dwError = GetLastError();
        if (dwError == ERROR_CANCELLED)
        {
            // 用户拒绝了UAC提示
            return FALSE;
        }
        return FALSE;
    }
    
    return TRUE;
}

// DLL入口点
BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}